package practice.tree.four;

import static practice.tree.four.RedBlackNode.BLACK;
import static practice.tree.four.RedBlackNode.RED;

/**
 * @Description:
 * @Author Ammar
 * @Create 2023/8/21 20:31
 */
public class RedBlackTree extends BalanceBinarySearchTree {
    @Override
    protected void afterAdd(TreeNode node) {
        TreeNode parent = node.parent;
        if (parent == null) {
            black(node);
            return;
        }

        if (isBlack(parent)) return;

        TreeNode uncle = parent.sibling();
        TreeNode grand = red(parent.parent);
        if (isRed(uncle)) {
            black(parent);
            black(uncle);
            afterAdd(grand);
            return;
        }

        if (parent.isLeftChild()) {
            if (node.isLeftChild()) {
                black(parent);
            } else {
                black(node);
                rotateLeft(parent);
            }
            rotateRight(grand);
        } else {
            if (node.isRightChild()) {
                black(parent);
            } else {
                black(node);
                rotateRight(parent);
            }
            rotateLeft(grand);
        }
    }

    @Override
    protected void afterRemove(TreeNode node, TreeNode replacement) {
        TreeNode parent = node.parent;
        if (isRed(node)) return;
        if (isRed(replacement)) {
            black(replacement);
            return;
        }
        if (parent == null) return;
        boolean left = parent.left == null || node.isLeftChild();
        TreeNode sibling = left ? parent.right : parent.left;
        if (left) {
            if (isRed(sibling)) {
                red(parent);
                black(sibling);
                rotateLeft(parent);
                sibling = parent.right;
            }

            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                boolean parentColor = colorOf(parent);
                black(parent);
                red(sibling);
                if (parentColor == BLACK) {
                    afterRemove(parent, null);
                }
            } else {
                if (isBlack(sibling.right)) {
                    rotateRight(sibling);
                    sibling = parent.right;
                }
                color(sibling, colorOf(parent));
                black(parent);
                black(sibling.right);
                rotateLeft(parent);
            }
        } else {
            if (isRed(sibling)) {
                red(parent);
                black(sibling);
                rotateRight(parent);
                sibling = parent.left;
            }

            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                boolean parentColor = colorOf(parent);
                black(parent);
                red(sibling);
                if (parentColor == BLACK) {
                    afterRemove(parent, null);
                }
            } else {
                if (isBlack(sibling.left)) {
                    rotateLeft(sibling);
                    sibling = parent.left;
                }
                color(sibling, colorOf(parent));
                black(parent);
                black(sibling.left);
                rotateRight(parent);
            }
        }
    }



    private TreeNode black(TreeNode node) {
        return color(node, BLACK);
    }

    private TreeNode red(TreeNode node) {
        return color(node, RED);
    }

    private TreeNode color(TreeNode node, boolean color) {
        if (node == null) return null;
        ((RedBlackNode) node).color = color;
        return node;
    }

    private boolean isBlack(TreeNode node) {
        return colorOf(node) == BLACK;
    }

    private boolean isRed(TreeNode node) {
        return colorOf(node) == RED;
    }

    private boolean colorOf(TreeNode node) {
        if (node == null) return BLACK;
        return ((RedBlackNode) node).color;
    }

    @Override
    protected TreeNode createTreeNode(Integer val, TreeNode parent) {
        return new RedBlackNode(val, parent);
    }


}
